#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File        : mngs_mjn_oumeng.py
@Author      : Bing Liang
@Email       : believer19940901@gmail.com
@Date        : 2025/11/18 10:52
@Description : 生成欧盟报告所需的文件，主要用于构建 shell 命令、
               执行任务、校验输出文件等辅助流程。
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path
import subprocess
import logging
from typing import Sequence
from multiprocessing import Pool
from datetime import datetime


# ----------------------------
# 日志配置
# ----------------------------
logging.basicConfig(
    level=logging.INFO,
    format="[%(asctime)s] - [%(levelname)s] - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
LOGGER = logging.getLogger("MNGS-MJN")


# ----------------------------
# 软件环境
# ----------------------------
PYTHON_ENV = Path("/home/bioinfo/software/miniconda3/envs/python-venv/bin")
BWA_ENV = Path("/home/bioinfo/software/miniconda3/envs/bwa/bin")
SAMTOOLS_ENV = Path("/home/bioinfo/software/miniconda3/envs/psamtools/bin")


# ----------------------------
# 脚本路径
# ----------------------------
SCRIPT_DIR = Path(__file__).parent
MNGS_MJN_SAMPLE_QSTAT_PY = SCRIPT_DIR / "mngs_mjn_sample_qstat.py"
MNGS_MJN_SAMPLE_AI_PY = SCRIPT_DIR / "mngs_mjn_sample_ai.py"
MNGS_MJN_MERGE_PY = SCRIPT_DIR / "mngs_mjn_merge.py"
MNGS_MJN_SELECT_SEQ_PY = SCRIPT_DIR / "mngs_mjn_select_seq.py"
MNGS_MJN_SELECT_PAN_PY = SCRIPT_DIR / "mngs_mjn_select_pan.py"
MNGS_MJN_PLOT_COVERAGE_PY = SCRIPT_DIR / "mngs_mjn_plot_coverage.py"


# ----------------------------
# 辅助函数
# ----------------------------
def _run_or_die(sh_cmd: str, sh_file: Path, log_file: Path, ok_file: Path, skip_ok: bool = True) -> None:
    """
    执行单个 shell 命令。若执行成功，生成 `.ok` 文件。
    若失败则记录日志，但不抛出异常（由外部统一检查）。

    Parameters
    ----------
    sh_cmd : str
        要写入并执行的 shell 命令内容。
    sh_file : Path
        shell 脚本文件路径。
    log_file : Path
        运行日志输出文件路径。
    ok_file : Path
        用于标记任务成功完成的 .ok 文件。
    """
    if ok_file.exists() and skip_ok:
        LOGGER.info(f"已跳过：{sh_file}（存在 .ok 文件）")
        return

    try:
        LOGGER.info(f"开始执行：{sh_file}")

        # 写入 shell 文件
        with open(sh_file, "w", encoding="utf-8") as fs:
            fs.write(sh_cmd)

        # 执行脚本并捕获输出
        with open(log_file, "w", encoding="utf-8") as fl:
            result = subprocess.run(
                f"bash {sh_file}",
                shell=True,
                text=True,
                stdout=fl,
                stderr=fl,
            )

        if result.returncode != 0:
            LOGGER.error(
                f"执行失败：{sh_file}，退出码={result.returncode}"
            )
            return

        # 任务成功，创建 OK 文件
        ok_file.touch()
        LOGGER.info(f"执行完成：{sh_file}")

    except Exception as e:
        LOGGER.error(f"执行异常：{sh_file}，错误：{e}")


def _run(
    sh_files: Sequence[Path],
    sh_cmds: Sequence[str],
    cpu_n: int = 1,
    check_ok: bool = True
) -> None:
    """
    并行执行多个 shell 脚本。

    Parameters
    ----------
    sh_files : Sequence[Path]
        要生成的 shell 脚本文件列表。
    sh_cmds : Sequence[str]
        shell 命令字符串列表。与 sh_files 位置一一对应。
    cpu_n : int, optional
        并行进程数，默认 1。
    check_ok : bool, optional
        是否在执行结束后检查所有 .ok 文件，默认 True。

    Raises
    ------
    RuntimeError
        若任务要求检查 ok 文件但某些 .ok 缺失，则抛出异常。
    """
    ok_files = [f.with_suffix(".ok") for f in sh_files]
    log_files = [f.with_suffix(".log") for f in sh_files]

    with Pool(cpu_n) as pool:
        for sh_cmd, sh_file, ok_file, log_file in zip(sh_cmds, sh_files, ok_files, log_files):
            pool.apply_async(
                _run_or_die,
                args=(sh_cmd, sh_file, log_file, ok_file)
            )
        pool.close()
        pool.join()

    if check_ok:
        failed = [sh for sh, ok in zip(sh_files, ok_files) if not ok.exists()]
        if failed:
            raise RuntimeError(
                "以下步骤未生成 OK 文件：\n  " + "\n  ".join(map(str, failed))
            )


def _check(out_files: Sequence[Path], empty_ok: bool = False) -> None:
    """
    检查指定输出文件是否存在，是否为空。

    Parameters
    ----------
    out_files : Sequence[Path]
        输出文件路径列表。
    empty_ok : bool, optional
        是否允许空文件，默认 False。

    Raises
    ------
    RuntimeError
        当文件不存在或为空（且 empty_ok=False）时抛出。
    """
    missing_files = [f for f in out_files if not f.exists()]
    empty_files = [f for f in out_files if f.exists() and f.stat().st_size == 0]

    if missing_files:
        raise RuntimeError(
            "以下文件不存在：\n  " + "\n  ".join(map(str, missing_files))
        )

    if empty_files and not empty_ok:
        raise RuntimeError(
            "以下文件为空：\n  " + "\n  ".join(map(str, empty_files))
        )


def _parse_args() -> Namespace:
    """
    解析命令行参数。

    本函数用于解析任务执行所需的基本输入信息，包括芯片编号与样本编号。
    返回的 Namespace 对象可通过属性访问方式获取参数值，例如：
        args.job_id
        args.sample_id

    Returns
    -------
    Namespace
        解析后的命令行参数对象。
    """
    parser = ArgumentParser(
        description="根据指定模板生成病原体报告"
    )

    parser.add_argument(
        "--job_id",
        type=str,
        required=True,
        help="芯片编号（任务或流程执行的唯一标识）"
    )

    parser.add_argument(
        "--sample_id",
        type=str,
        required=True,
        help="样本编号（用于定位当前样本的数据与分析结果）"
    )

    return parser.parse_args()


# --------
# 分析流程
# --------
class Begin:
    """
    流程启动 - 获取命令行参数
    """

    def __init__(self, args: Namespace) -> None:
        self.step_n = 0
        self.begin_time = datetime.now()
        self.job_id = args.job_id.strip()
        self.sample_id = args.sample_id.strip()
        self.out_dir = Path("/data/mNGS/runmngs/MaiJN") / self.job_id / self.sample_id
        self.log_dir = self.out_dir / "log"


    def run(self) -> None:
        LOGGER.info(f"TASK - {self.job_id} - {self.sample_id} - STARTED")
        self.out_dir.mkdir(parents=True, exist_ok=True)
        self.log_dir.mkdir(parents=True, exist_ok=True)


class SampleQstat(Begin):
    """
    生成样本质控结果
    """
    def __init__(self, args: Namespace) -> None:
        super().__init__(args)
        self.step_n += 1
        self.sample_qstat_out_name = f"step{self.step_n}_sample_qstat"
        self.sample_qstat_out_dir = self.out_dir / self.sample_qstat_out_name
        self.sample_qstat_log_dir = self.log_dir / self.sample_qstat_out_name
        self.sample_qstat_out_file = self.sample_qstat_out_dir / f"{self.sample_qstat_out_name}.tsv"

    def run(self) -> None:
        super().run()
        self.sample_qstat_log_dir.mkdir(parents=True, exist_ok=True)
        self.sample_qstat_log_dir.mkdir(parents=True, exist_ok=True)
        _run(
            [
                self.sample_qstat_log_dir / f"{self.sample_qstat_out_name}.sh"
            ],
            [
                (
                    f"#!/bin/bash\n\n"
                    f"set -vexuo pipefail\n\n"
                    f"export PATH={PYTHON_ENV}:$PATH\n\n"
                    f"python3 {MNGS_MJN_SAMPLE_QSTAT_PY} "
                    f"--sample_id {self.sample_id} "
                    f"--job_id {self.job_id} "
                    f"--out_file {self.sample_qstat_out_file}\n\n"
                )
            ],
            1
        )
        _check([self.sample_qstat_out_file])


class SampleAi(SampleQstat):
    """
    生成欧蒙检出结果
    """
    def __init__(self, args: Namespace) -> None:
        super().__init__(args)
        self.step_n += 1
        self.sample_ai_out_name = f"step{self.step_n}_sample_ai"
        self.sample_ai_out_dir = self.out_dir / self.sample_ai_out_name
        self.sample_ai_log_dir = self.log_dir / self.sample_ai_out_name
        self.sample_ai_out_file = self.sample_ai_out_dir / f"{self.sample_ai_out_name}.tsv"
        self.sample_ai_out_rgi = self.sample_ai_out_dir / f"{self.sample_ai_out_name}.rgi.tsv"
        self.sample_ai_out_vf = self.sample_ai_out_dir / f"{self.sample_ai_out_name}.vf.tsv"

    def run(self) -> None:
        super().run()
        self.sample_ai_log_dir.mkdir(parents=True, exist_ok=True)
        self.sample_ai_out_dir.mkdir(parents=True, exist_ok=True)
        _run(
            [
                self.sample_ai_log_dir / f"{self.sample_ai_out_name}.sh"
            ],
            [
                (
                    f"#!/bin/bash\n\n"
                    f"set -vexuo pipefail\n\n"
                    f"export PATH={PYTHON_ENV}:$PATH\n\n"
                    f"python3 {MNGS_MJN_SAMPLE_AI_PY} "
                    f"--job_id {self.job_id} "
                    f"--sample_id {self.sample_id} "
                    f"--out_rgi {self.sample_ai_out_rgi} "
                    f"--out_vf {self.sample_ai_out_vf} "
                    f"--out_file {self.sample_ai_out_file}\n\n"
                )
            ],
            1
        )
        _check([self.sample_ai_out_file])


class Merge(SampleAi):
    """
    合并所有tsv文件到excel文件中
    """
    def __init__(self, args: Namespace) -> None:
        super().__init__(args)
        self.step_n += 1
        self.merge_out_name = f"step{self.step_n}_merge"
        self.merge_out_dir = self.out_dir / self.merge_out_name
        self.merge_log_dir = self.log_dir / self.merge_out_name
        self.merge_out_file = self.merge_out_dir / f"{self.merge_out_name}.xlsx"

    def run(self) -> None:
        super().run()
        self.merge_log_dir.mkdir(parents=True, exist_ok=True)
        self.merge_out_dir.mkdir(parents=True, exist_ok=True)
        _run(
            [self.merge_log_dir / f"{self.merge_out_name}.sh"],
            [
                (
                    f"#!/bin/bash\n\n"
                    f"set -vexuo pipefail\n\n"
                    f"export PATH={PYTHON_ENV}:$PATH\n\n"
                    f"python3 {MNGS_MJN_MERGE_PY} "
                    f"--tsv1 {self.sample_qstat_out_file} "
                    f"--tsv2 {self.sample_ai_out_file} "
                    f"--tsv3 {self.sample_ai_out_rgi} "
                    f"--tsv4 {self.sample_ai_out_vf} "
                    f"--out_xlsx {self.merge_out_file}\n\n"
                )
            ],
            1
        )
        _check([self.merge_out_file])


class PlotCoverage(Merge):
    """
    绘制覆盖率图
    """
    def __init__(self, args: Namespace) -> None:
        super().__init__(args)
        self.step_n += 1
        self.plot_coverage_out_name = f"step{self.step_n}_plot_coverage"
        self.plot_coverage_out_dir = self.out_dir / self.plot_coverage_out_name
        self.plot_coverage_log_dir = self.log_dir / self.plot_coverage_out_name

    def run(self) -> None:
        super().run()
        self.plot_coverage_log_dir.mkdir(parents=True, exist_ok=True)
        self.plot_coverage_out_dir.mkdir(parents=True, exist_ok=True)
        _run(
            [   self.plot_coverage_log_dir / f"{self.plot_coverage_out_name}.sh"
            ],
            [
                (
                    f"#!/bin/bash\n\n"
                    f"set -vexuo pipefail\n\n"
                    f"export PATH={PYTHON_ENV}:$PATH\n\n"
                    f"python3 {MNGS_MJN_PLOT_COVERAGE_PY} "
                    f"--job_id {self.job_id} "
                    f"--sample_id {self.sample_id} "
                    f"--in_tsv {self.sample_ai_out_file} "
                    f"--out_dir {self.plot_coverage_out_dir}\n\n"
                )
            ],
            1
        )


class End(PlotCoverage):
    """最终步骤：输出耗时"""

    def __init__(self, args: Namespace):
        super().__init__(args)

    def run(self):
        super().run()
        elapsed = datetime.now() - self.begin_time
        LOGGER.info(f"任务完成，总耗时：{elapsed}")


if __name__ == "__main__":
    End(_parse_args()).run()